
import warnings
warnings.simplefilter("ignore")
import holoviews as hv
from holoviews.operation.datashader import rasterize
import geoviews as gv
from geoviews.tile_sources import EsriImagery
hv.extension('bokeh')
from msatutil.msat_dset import msat_dset # https://github.com/rocheseb/msatutil


from IPython.display import display, HTML
display(HTML("<style>.container { width:100% !important; }</style>"))


def show_map(x,y,z,width=450,height=450,cmap="viridis",clim=None,title="",background_tile=EsriImagery):
    """
    Make a geoviews map of z overlayed on background_tile
    This doesn't preserve pixel shapes as the image is re-rasterized after each zoom/pan
    
    Inputs:
        x: 1D or 2D array of longitudes (shape (N,) or (M,N))
        y: 1D or 2D array of latitudes (shape (M,) or (M,N))
        z: 2D array of the data to plot (shape (M,N))
        width (int): plot width in pixels
        height (int): plot height in pixels
        cmap (str): named colormap
        clim (Tuple(float,float)): z-limits for the colorbar, give False to use dynamic colorbar
        title (str): plot title
        background_tile: the geoviews tile the plot will be made over and that will be in the linked 2nd panel
        
    Outputs:
        geoviews figure
    """
    
    quad = gv.project(gv.QuadMesh((x,y,z)))
    
    if clim is None:
        # define color limits as mean +/- 3 std
        mean_z = np.nanmean(z)
        std_z = np.nanstd(z,ddof=1)
        clim = (mean_z-3*std_z,mean_z+3*std_z)

    raster = rasterize(
        quad,
        width=width,
        height=height
    ).opts(
        width=width,
        height=height,
        cmap=cmap,
        colorbar=True,
        title=title,
    )
    
    if clim is not False:
        raster = raster.opts(clim=clim)
    
    # Make a dummy quadmesh that will have alpha=0 in the second panel so we can see the EsriImagery under the data
    # I do this so it will add a colorbar on the second plot so we don't need to think about resizing it
    # just use a small subset of data so it doesn't trigger much computations
    if x.ndim == 1:
        xdum = x[:10]
        ydum = y[:10]
    else:
        xdum = x[:10,:10]
        ydum = y[:10,:10]
    dummy = gv.project(gv.QuadMesh((xdum,ydum,z[:10,:10]))).opts(
        width=width,
        height=height,
        cmap=cmap,
        colorbar=True,
        alpha=0,
        title="Esri Imagery"
    )
    if clim is not False:
        dummy.opts(clim=clim)
    
    return background_tile * (raster + dummy)


l3 = msat_dset("gs://msat-prod-data-methaneair-level3/MAIR/2021/08/06/RF06/po-211/level3/5x1/mosaic/30m/MethaneAIR_L3_mosaic_20210806T161712_20210806T183039_dpp.nc")


show_map(l3["lon"][:],l3["lat"][:],l3["xch4"][:],title="RF06 L3 XCH4 (ppb)",width=850,height=750)